library(SingleCellExperiment)
library(scater)
library(scMerge)
library(tidyverse)

## Subsetted mouse ESC data
data("example_sce", package = "scMerge")

example_sce = runPCA(example_sce, exprs_values = "logcounts")

plotPCA(
  example_sce, 
  colour_by = "cellTypes", 
  shape_by = "batch")

# 1. obtaining negative ctrl: SEG ---------
## single-cell stably expressed gene list
data("segList_ensemblGeneID", package = "scMerge")
head(segList_ensemblGeneID$mouse$mouse_scSEG)

# 2. unsupervised scMerge ----------
# kmeansK value means there are 3 types of cell in each sample
scMerge_unsupervised <- scMerge(
  sce_combine = example_sce, 
  ctl = segList_ensemblGeneID$mouse$mouse_scSEG,
  kmeansK = c(3, 3),
  assay_name = "scMerge_unsupervised"
  )

scMerge_unsupervised = runPCA(scMerge_unsupervised, exprs_values = "scMerge_unsupervised")
plotPCA(
  scMerge_unsupervised, 
  colour_by = "cellTypes", 
  shape_by = "batch")

scMerge_unsupervised@assays@data$counts[1:5,1:5]
scMerge_unsupervised@assays@data$logcounts[1:5,1:5]
scMerge_unsupervised@assays@data$scMerge_unsupervised[1:5,1:5]

convert_neg2zero <- function(x){
  ifelse(x<0,0,x)
}

convert_neg2zero(-3:3)

nonge_mat <- convert_neg2zero(adjusted_mat)

# out of RAM computation ----
library(HDF5Array)
library(DelayedArray)

DelayedArray:::set_verbose_block_processing(TRUE) ## To monitor block processing 

hdf5_input = example_sce

assay(hdf5_input, "counts") = as(counts(hdf5_input), "HDF5Array")
assay(hdf5_input, "logcounts") = as(logcounts(hdf5_input), "HDF5Array")
system.time(scMerge_hdf5 <- scMerge(
  sce_combine = hdf5_input, 
  ctl = segList_ensemblGeneID$mouse$mouse_scSEG,
  kmeansK = c(3, 3),
  assay_name = "scMerge_hdf5",
  BSPARAM = BiocSingular::RandomParam(),
  BPPARAM = BiocParallel::MulticoreParam(workers = 6),
  verbose = TRUE)) 

# hierarchical scMerge2 -----------
# Create a fake sample information
example_sce$sample <- rep(c(1:4), each = 50)
table(example_sce$sample, example_sce$batch)

# Construct a hierarchical index list
# level1 in batch among samples, level2 between batches
h_idx_list <- list(level1 = split(seq_len(ncol(example_sce)), example_sce$batch),
                   level2 = list(seq_len(ncol(example_sce))))

# Construct a batch information list
batch_list <- list(level1 = split(example_sce$sample, example_sce$batch),
                   level2 = list(example_sce$batch))

# need to set ruvK for each level, low level with lower K value
scMerge2_res <- scMerge2h(exprsMat = logcounts(example_sce),
                          batch_list = batch_list,
                          h_idx_list = h_idx_list,
                          ctl = segList_ensemblGeneID$mouse$mouse_scSEG,
                          ruvK_list = c(2, 5))

# the last item in list is output of last level scMerge2h
assay(example_sce, "scMerge2h") <- last(scMerge2_res)
